
#if !defined __RTCN_RUDP_H__
#define __RTCN_RUDP_H__

#include "RtCnRudpBase.h"
#include "RtConnBase.h"
#include "RtCnBase.h"
#include "BitStream.h"
#include "SingleProducerConsumer.h"
#include "DS_RangeList.h"
#include "RtQueueT.h"
#include "DS_LinkedList.h"
#include "DS_BPlusTree.h"
#include "RtCnRudpPacketPool.h"

class CRtRudpFlowControl;
class CRtRudpConnClient;
class CRtRudpConnServer;
class IRtObserver;

typedef ServerListT<CRtRudpConnServer>			CRtRudpConnServerList;
typedef CConnAcceptorSinkT<CRtRudpConnServer>	CRtRudpConnAcceptorSink;
typedef CConnAcceptorT<CRtRudpConnServer>		CRtRudpConnAcceptor;
typedef CConnConnectorT<CRtRudpConnClient>		CRtRudpConnConnector;

/// NOTE: Should not use std::map as RtRudpOrderingID may wrap. 
typedef	DataStructures::LinkedList<RtRudpPacket*> RudpRecvOrderingListType;

int SplitPacketIndexComp( RtRudpSplitPacketID  &key, RtRudpPacket*  &data );

struct SplitPacketChannel
{
	RtRudpTimeUS lastUpdateTime;
	DataStructures::OrderedList<RtRudpSplitPacketID, RtRudpPacket*, SplitPacketIndexComp> splitPacketList;
};

int SplitPacketChannelComp( RtRudpSplitPacketNo  &key, SplitPacketChannel*  &data );

class CRtRudpConn: public CRtConnBase
{
	friend class CRtRudpFlowControl;
	friend class CRtRudpSendBufferedEvent;
public:
	virtual RtResult SendData(
		CRtMessageBlock& inData,
		BOOL inSendImmediately = TRUE,
		BOOL inDestroyData = FALSE,
		DataPriority inPriority = DP_MEDIUM, 
		BOOL inNeedAck = TRUE);

	virtual RtResult SendData(
		CRtMessageBlock& inData, 
		RtRudpReliability inReliability,
		INT inTTL = -1,
		BOOL inSendImmediately = TRUE,
		BOOL inDestroyData = FALSE,
		DataPriority inPriority = DP_MEDIUM,
		RudpChannel inChannelID = 0); 

	virtual RtResult SetOption(
		DWORD aCommand, 
		LPVOID aArg);

	virtual RtResult GetOption(
		DWORD aCommand, 
		LPVOID aArg);

	virtual RtResult Disconnect(
		RtResult aReason);

	virtual void OnSend(
		IRtTransport *aTrptId,
		CRtTransportParameter *aPara = NULL);

	virtual void OnTimer(CRtTimerWrapper* aId);

public:
	void NeedKeepAlive(BOOL bNeedKeepAlive) { m_needKeepAlive = bNeedKeepAlive; }
	const CRtInetAddr& GetLocalAddr() const { return m_localAddr; }
	const CRtInetAddr& GetPeerAddr()const { return m_peerAddr; }
	RtResult RunSendCycle(RtRudpTimeUS time);
#ifdef RUDP_FLOWCONTROL_OPTIMIZE
	BOOL InsertIntoFlowControl();
	BOOL RemoveFromFlowControl();
	BOOL NeedFlowControl();
#endif

public:
	CRtRudpConn();
	virtual ~CRtRudpConn();

protected:
	void InitializeVariables( void );
	void SetPing( RtRudpTime i );
	void OnReceive_i(
		CRtMessageBlock &aData,
		IRtTransport *aTrptId,
		CRtTransportParameter *aPara = NULL);

	/// Puts data on the m_sendPduBuffer
	RtResult SendBuffered(
		RtRudpPduType aType,
		CRtMessageBlock &aData, 
		RtRudpReliability aReliability,
		DataPriority aPriority = DP_MEDIUM,
		RudpChannel aChannelID = 0,
		INT aTTL = -1,
		BOOL inDestroyData = FALSE); 

	RtResult SendBuffered_i();

	/// Puts data on the sendPacketSet
	RtResult SendPdu(
		RtRudpPduType aType,
		CRtMessageBlock &aData, 
		RtRudpReliability aReliability,
		DataPriority aPriority = DP_MEDIUM,
		RudpChannel aChannelID = 0,
		INT aTTL = -1,
		BOOL aDuplicateData = TRUE); 

	/// Split the passed packet into chunks under MTU_SIZE bytes (including headers) and save those new chunks
	void SplitPacket( RtRudpPacket *internalPacket );

	/// Generates a datagram (coalesced packets)
	/// \param[out] output The resulting BitStream
	/// \param[in] Current MTU size
	/// \param[out] reliableDataSent Set to true or false as a return value as to if reliable data was sent.
	/// \param[in] time Current time
	/// \return The number of messages sent
	unsigned GenerateDatagram( RakNet::BitStream *output, int& reliableDataSentNum, RtRudpTimeUS time, BOOL updateResendQueue = TRUE );

	/// Parse an internalPacket and figure out how many header bits would be written.  Returns that number
	int GetBitStreamHeaderLength( const RtRudpPacket *const internalPacket );

	///Parse an internalPacket and create a bitstream to represent this data. Returns number of bits used
	int WriteToBitStreamFromInternalPacket( RakNet::BitStream *bitStream, const RtRudpPacket *const internalPacket );

	/// Inserts a packet into the resend list in order
	void InsertPacketIntoResendList( RtRudpPacket *internalPacket, RtRudpTimeUS time, bool firstResend );

	/// Send the contents of a bitstream to the socket
	/// \param[in] bitStream The data to send.
	RtResult SendBitStream( RakNet::BitStream *bitStream );

	/// Does what the function name says
	unsigned RemovePacketFromResendList(const RtRudpPacketNo packetNo, RtRudpTimeUS time, RtRudpTime& rtt);

	/// Parse a bitstream and create an internal packet to represent this data
	RtRudpPacket* CreateInternalPacketFromBitStream( RakNet::BitStream *bitStream, RtRudpTimeUS time );

	/// Acknowledge receipt of the packet with the specified messageNumber
	void SendAcknowledgementPacket( const RtRudpPacketNo messageNumber, RtRudpTimeUS time ) { acknowlegements.Insert(messageNumber); }

	/// Delete any unreliable split packets that have long since expired
	void DeleteOldUnreliableSplitPackets( RtRudpTimeUS time );

	/// Returns true if newPacketOrderingIndex is older than the waitingForPacketOrderingIndex
	bool IsOlderOrderedPacket( RtRudpOrderingID newPacketOrderingIndex, RtRudpOrderingID waitingForPacketOrderingIndex );

	/// Insert a packet into the split packet list
	void InsertIntoSplitPacketList( RtRudpPacket * internalPacket, RtRudpTimeUS time );

	/// Take all split chunks with the specified splitPacketId and try to reconstruct a packet. If we can, allocate and return it.  Otherwise return 0
	RtRudpPacket * BuildPacketFromSplitPacketList( RtRudpSplitPacketNo splitPacketId, RtRudpTimeUS time );

	/// Get the specified ordering list
	RudpRecvOrderingListType *GetOrderingListAtOrderingStream( unsigned char orderingChannel );

	/// Add the internal packet to the ordering list in order based on order index
	void AddToOrderingList( RtRudpPacket * internalPacket );

	RtResult UpdateSending(RtRudpTimeUS time, BOOL updateResendQueue = TRUE);

protected:
	void HandleReceivedPacket(RtRudpPacket* aPacket);
	void HandleShutdownPdu();
	void HandleDataPdu(CRtMessageBlock& aData);
#if 0
	void HandleDropPacket(RtRudpPacket* aPacket, RtRudpTimeUS inCurTime);
#endif
	void ReceivedPacketQueuePopExpiredTimeAndCompress(RtRudpTimeUS time);
	void DeleteOldSplitPackets(RtRudpSplitPacketNo splitPacketId);
	void ClearSendPduBuffer();

	/// Move all PDUs from m_sendPduBuffer to sendPacketSet.
	void WrapUpPdusToPackets(/*RtRudpTimeUS time*/); 

	RtResult SendData_i(DWORD aLength, LPBYTE aData);
	void DecreaseSendBufferSize(const RtRudpPacket& aPacket);
	void CheckOnSend();

	void Release();

protected:
	virtual void Close_i();
	virtual void HandleConnReqPdu(CRtMessageBlock& aData) { RT_ASSERTE(FALSE); }
	virtual void HandleConnRespPdu(CRtMessageBlock& aData) { RT_ASSERTE(FALSE); }
	virtual void CacheReceivedData(CRtMessageBlock& aData) { RT_ASSERTE(FALSE); }

protected:
	CRtRudpFlowControl* m_pFlowControl;

	CRtThread* m_pThreadNetwork;

	CRtInetAddr m_localAddr;
	CRtInetAddr m_peerAddr;

	INT m_mtuSize;
//	RtRudpTimeUS m_lastReliableSendTime;
	RtRudpTimeUS m_lastPacketReceiveTime;
	DataStructures::SingleProducerConsumer<RtRudpPdu> m_sendPduBuffer;
	DWORD m_maxSendBufferSize;
	DWORD m_totalSizeOfPassedInData;
	DWORD m_totalSizeOfSentOutData;
#define RUDP_SEND_BUFFER_SIZE (m_totalSizeOfPassedInData-m_totalSizeOfSentOutData)
	BOOL m_needKeepAlive;
	BOOL m_needOnSend;
#ifdef RUDP_FLOWCONTROL_OPTIMIZE
	BOOL m_onFlowControl;
#endif
	BOOL m_isAckSent; // If any ACK has been sent out in the last RunSendCycle. This member is added for performance.
	BOOL m_isDataSent; // If any data has been sent out in recent keep-alive interval period.
	WORD m_recvDropPackets;
	WORD m_sendDropPackets;
	WORD m_recvDropTick;
	WORD m_sendDropTick;
	IRtObserver* m_observer;
	WORD m_resendIncrement;
//	RtRudpTimeUS m_lastWarningResendQueueTime;

#ifdef RUDP_ENABLE_STATISTICS
	CRtRudpStatistics m_statistics;
#endif

protected:
	DataStructures::List<RudpRecvOrderingListType*> orderingList;
	DataStructures::RangeList<RtRudpPacketNo> acknowlegements;

	DataStructures::BPlusTree<RtRudpPacketNo, RtRudpPacket*, RESEND_TREE_ORDER> resendList;
	DataStructures::Queue<RtRudpPacket*> resendQueue;

	DataStructures::Queue<RtRudpPacket*> sendPacketSet[ DP_NUMBER_OF_PRIORITIES ];
	DataStructures::OrderedList<RtRudpSplitPacketNo, SplitPacketChannel*, SplitPacketChannelComp> splitPacketChannelList;
	RtRudpPacketNo messageNumber;
	RtRudpTimeUS lastAckTime;
	RakNet::BitStream updateBitStream;
	RtRudpOrderingID waitingForOrderedPacketWriteIndex[ NUMBER_OF_ORDERED_STREAMS ], waitingForSequencedPacketWriteIndex[ NUMBER_OF_ORDERED_STREAMS ];
	RtRudpOrderingID waitingForOrderedPacketReadIndex[ NUMBER_OF_ORDERED_STREAMS ], waitingForSequencedPacketReadIndex[ NUMBER_OF_ORDERED_STREAMS ];
	
	RtRudpTime ping;
	RtRudpSplitPacketNo splitPacketId;
	RtRudpTimeUS timeoutTime; // How long to wait before timing someone out

	/// Memory-efficient receivedPackets algorithm:
	/// receivedPacketsBaseIndex is the packet number we are expecting
	/// Everything under receivedPacketsBaseIndex is a packet we already got
	/// Everything over receivedPacketsBaseIndex is stored in hasReceivedPacketQueue
	/// It stores the time to stop waiting for a particular packet number, where the packet number is receivedPacketsBaseIndex + the index into the queue
	/// If 0, we got got that packet.  Otherwise, the time to give up waiting for that packet.
	/// If we get a packet number where (receivedPacketsBaseIndex-packetNumber) is less than half the range of receivedPacketsBaseIndex then it is a duplicate
	/// Otherwise, it is a duplicate packet (and ignore it).
	DataStructures::Queue<RtRudpTimeUS> hasReceivedPacketQueue;
	RtRudpPacketNo receivedPacketsBaseIndex;

	RtRudpTimeUS lastUpdateTime;

	double availableBandwidth, currentBandwidth; // In bits per second. 
#ifdef RUDP_ENABLE_BW_DETECT
	RtRudpTimeUS histogramEndTime, histogramStartTime;
	unsigned histogramReceiveMarker;
	int noPacketlossIncreaseCount;
	unsigned histogramPlossCount, histogramAckCount;
	double lowBandwidth, highBandwidth;
	bool continuousSend;
#endif // RUDP_ENABLE_BW_DETECT

	RtRudpTimeUS resendTimeIncrement;
};

class CRtRudpConnClient: public CRtRudpConn
{
public:
	virtual RtResult Disconnect(RtResult inReason);

	virtual void OnReceive(
		CRtMessageBlock &aData,
		IRtTransport *aTrptId,
		CRtTransportParameter *aPara = NULL);

	virtual void OnDisconnect(
		RtResult aReason,
		IRtTransport *aTrptId);
	
	virtual void OnConnectIndication(
		RtResult aReason,
		IRtTransport *aTrpt,
		IRtAcceptorConnectorId *aRequestId);

	virtual void OnTimer(CRtTimerWrapper* aId);

public:
	CRtRudpConnClient(CRtConnectionManager::CType);
	virtual ~CRtRudpConnClient();

	void SetConnConnector(CRtRudpConnConnector* pConnConnector) { m_pConnConnector = pConnConnector; }
	void CancelHandShake();
	void ScheduleConnectTimer(const CRtTimeValue& aTimeout);

protected:
	RtResult SendSync();
	virtual void Close_i();
	virtual void HandleConnRespPdu(CRtMessageBlock& aData);
	virtual void CacheReceivedData(CRtMessageBlock& aData);
	void ClearReceivedData(IRtTransportSink* aTransportSink = NULL);
	DWORD GetLocalIP();

protected:
	CRtRudpConnConnector* m_pConnConnector;
	BYTE m_sendSyncCount;
	/// m_receivedData is used to cache the incoming upper-layer data of client if it is, probably, received  
	/// before CONNRESP pdu when the client is connecting to the server.
	list<CRtMessageBlock*> m_receivedData;

	CRtTimerWrapper m_sendSyncTimer;
	RtRudpTime m_lastPing;
};

class CRtRudpConnServer: public CRtRudpConn
{
public:
	virtual void OnReceive(
		CRtMessageBlock &aData,
		IRtTransport *aTrptId,
		CRtTransportParameter *aPara = NULL);

	virtual void OnDisconnect(
		RtResult aReason,
		IRtTransport *aTrptId);
	
	virtual void OnConnectIndication(
		RtResult aReason,
		IRtTransport *aTrpt,
		IRtAcceptorConnectorId *aRequestId);

public:
	CRtRudpConnServer(CRtConnectionManager::CType);
	virtual ~CRtRudpConnServer();

	void SetConnAcceptor(CRtRudpConnAcceptor* pConnAcceptor) { m_pConnAcceptor = pConnAcceptor; }
	
protected:
	virtual void HandleConnReqPdu(CRtMessageBlock& aData);

protected:
	CRtAutoPtr<CRtRudpConnAcceptor> m_pConnAcceptor;
	CRtInetAddr m_peerPrivateAddr;
};

class CRtRudpSendBufferedEvent: public IRtEvent
{
public:
	CRtRudpSendBufferedEvent(CRtRudpConn* aConn)
		: m_pConn(aConn)
	{}
	~CRtRudpSendBufferedEvent() {}

	virtual RtResult OnEventFire() 
	{
		if (m_pConn->GetStatus()==RUDP_STATE_OPEN)
		{
			return m_pConn->SendBuffered_i();
		}
		return RT_OK;
	}

private:
	CRtRudpConn* m_pConn; // Here we don't need CRtAutoPtr as CRtRudpConn inherits from CRtReferenceControlSingleThreadTimerDelete
};

#ifdef RUDP_FLOWCONTROL_OPTIMIZE
class CRtRudpInsertIntoFlowControlEvent: public IRtEvent
{
public:
	CRtRudpInsertIntoFlowControlEvent(CRtRudpConn* aConn)
		: m_pConn(aConn)
	{}
	~CRtRudpInsertIntoFlowControlEvent() {}

	virtual RtResult OnEventFire() 
	{
		if (m_pConn->GetCurrStatus()==RUDP_STATE_OPEN)
		{
			m_pConn->InsertIntoFlowControl();
		}
		return RT_OK;
	}

private:
	CRtRudpConn* m_pConn;
};

inline BOOL CRtRudpConn::NeedFlowControl()
{
	if (m_wStatus==RUDP_STATE_CLOSED)
		return FALSE;

	// Need flow control if one of the following buffer is not empty:
	// resenList, acknowledgements, sendPduBuffer, sendPacketSet, updateBitStream

	if (!resendList.IsEmpty() || acknowlegements.Size()>0 || m_sendPduBuffer.Size()>0 || updateBitStream.GetNumberOfBitsUsed()>0)
		return TRUE;

	for (unsigned i = 0; i < RUDP_NUMBER_OF_PRIORITIES; ++i)
	{
		if (!sendPacketSet[i].IsEmpty())
			return TRUE;
	}

	return FALSE;
}
#endif // RUDP_FLOWCONTROL_OPTIMIZE

inline RtResult CRtRudpConn::SendData(
		CRtMessageBlock& inData,
		BOOL inSendImmediately,
		BOOL inDestroyData,
		DataPriority inPriority, 
		BOOL inNeedAck)
{
	return SendBuffered(RUDP_PDU_DATA, inData, RUDP_RELIABLE_ORDERED, inPriority, 0, -1, inDestroyData);
}

inline RtResult CRtRudpConn::SendData(
		CRtMessageBlock& inData, 
		RtRudpReliability inReliability,
		INT inTTL,
		BOOL inSendImmediately,
		BOOL inDestroyData,
		DataPriority inPriority,
		RudpChannel inChannelID)
{
	return SendBuffered(RUDP_PDU_DATA, inData, inReliability, inPriority, inChannelID, inTTL, inDestroyData);
}

inline RtResult CRtRudpConn::SendBuffered_i()
{
	WrapUpPdusToPackets();
	return UpdateSending(lastUpdateTime, FALSE);
}

inline RudpRecvOrderingListType* CRtRudpConn::GetOrderingListAtOrderingStream( unsigned char orderingChannel )
{
	if ( orderingChannel >= orderingList.Size() )
		return 0;

	return orderingList[ orderingChannel ];
}

inline void CRtRudpConn::WrapUpPdusToPackets(/*RtRudpTimeUS time*/)
{
	RtRudpPdu* pdu = NULL;
	while ( ( pdu=m_sendPduBuffer.ReadLock() )!=NULL )
	{
		SendPdu(pdu->pduType, *pdu->data, pdu->reliability, pdu->priority, pdu->orderingChannel, pdu->ttl, FALSE);
		m_sendPduBuffer.ReadUnlock();
		m_isDataSent = TRUE;
	}
}

inline RtResult CRtRudpConn::SendData_i(DWORD aLength, LPBYTE aData)
{
	if (m_lowerTransport==NULL)
		return RT_ERROR_NULL_POINTER;

	CRtMessageBlock mb(aLength, (LPCSTR)aData, CRtMessageBlock::DONT_DELETE, aLength);
	return m_lowerTransport->SendData(mb);
}

inline void CRtRudpConn::DecreaseSendBufferSize(const RtRudpPacket& aPacket)
{
	if (aPacket.pduType==RUDP_PDU_DATA && aPacket.data)
	{
		m_totalSizeOfSentOutData += aPacket.data->GetChainedLength();
	}
}

inline void CRtRudpConn::CheckOnSend()
{
	if (m_needOnSend && RUDP_SEND_BUFFER_SIZE<m_maxSendBufferSize && m_sink)
	{
		m_needOnSend = FALSE;
		m_sink->OnSend(this);
	}
}

inline DWORD CRtRudpConnClient::GetLocalIP()
{
	if (m_localAddr.GetIpAddrIn4Bytes()!=INADDR_ANY)
		return m_localAddr.GetIpAddrIn4Bytes();

	return ::RtGetLocalIP();
}

#endif	//__RTCN_RUDP_H__

